from fastapi import Request, Response
import time
from datetime import datetime
import json
from typing import Coroutine
import logging
import traceback

async def log_middleware(request: Request, call_next) -> Response:
    start_time = time.time()
    request_body = await request.body()
    
    try:
        response = await call_next(request)
    except Exception as e:
        process_time = time.time() - start_time
        error_detail = traceback.format_exc()
        
        logging.getLogger("error").error(json.dumps({
            "timestamp": datetime.now().isoformat(),
            "method": request.method,
            "path": str(request.url),
            "client_ip": request.client.host,
            "process_time": f"{process_time:.3f}s",
            "status_code": 500,
            "body": request_body.decode(),
            "error": str(e),
            "stack_trace": error_detail 
        }, ensure_ascii=False))
        raise
 
    process_time = time.time() - start_time
    
    # 获取响应体
    response_body = b""
    async for chunk in response.body_iterator:
        response_body += chunk
    response = Response(content=response_body, 
                        status_code=response.status_code,
                        headers=dict(response.headers),
                        media_type=response.media_type)

    log_data = {
        "timestamp": datetime.now().isoformat(),
        "method": request.method,
        "path": str(request.url),
        "client_ip": request.client.host,
        "process_time": f"{process_time:.3f}s",
        "status_code": response.status_code,
        "body": request_body.decode(),
        "response": response_body.decode()
    }
    
    # 记录API日志
    logging.getLogger("api").info(json.dumps(log_data, ensure_ascii=False))
    
    # 记录错误日志（状态码 >=400）
    if response.status_code >= 400:
        logging.getLogger("error").error(json.dumps(log_data, ensure_ascii=False))
    
    return response
    